// COLMAP - Structure-from-Motion and Multi-View Stereo.
// Copyright (C) 2016  Johannes L. Schoenberger <jsch at inf.ethz.ch>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

#include <Eigen/Geometry>
#include "base/polynomial.h"
#include "estimators/utils.h"
#include "p3p_Ke.h"
#include "util/types.h"
#include "util/poly.h"

namespace RVG
{

bool computePoses(const std::vector<P3PEstimator_KE::X_t> & bearing_vectors, const std::vector<P3PEstimator_KE::Y_t> & X_observations, std::vector<std::tuple<Mat3, Vec3>> & rotation_translation_solutions)
{
    const Vec3 w1 = X_observations.at(0);
    const Vec3 w2 = X_observations.at(1);
    const Vec3 w3 = X_observations.at(2);

    // k1
    const Vec3 u0 = w1 - w2;
    const double nu0 = u0.norm();
    const Vec3 k1 = u0.normalized();

    // bi
    Vec3 b1(bearing_vectors[0](0), bearing_vectors[0](1), 1.0);
    Vec3 b2(bearing_vectors[1](0), bearing_vectors[1](1), 1.0);
    Vec3 b3(bearing_vectors[2](0), bearing_vectors[2](1), 1.0);

    b1 = b1.normalized();
    b2 = b2.normalized();
    b3 = b3.normalized();

    // k3, tz
    Vec3 k3 = b1.cross(b2);
    const double nk3 = k3.norm();
    k3 = k3.normalized();

    const Vec3 tz = b1.cross(k3);

    // ui, vi
    const Vec3 v1 = b1.cross(b3);
    const Vec3 v2 = b2.cross(b3);

    const Vec3 u1 = w1 - w3;

    // coefficients related terms
    const double u1k1 = u1.dot(k1);
    const double k3b3 = k3.dot(b3);

    // f1i
    double f11 = k3b3;
    double f13 = k3.dot(v1);
    const double f15 = -u1k1 * f11;

    //delta
    const Vec3 nl = u1.cross(k1).normalized();
    const double delta = u1.cross(k1).norm();
    f11 *= delta;
    f13 *= delta;

    // f2i
    const double u2k1 = u1k1 - nu0;
    double f21 = tz.dot(v2);
    double f22 = nk3 * k3b3;
    double f23 = k3.dot(v2);
    const double f24 = u2k1 * f22;
    const double f25 = -u2k1 * f21;
    f21 *= delta;
    f22 *= delta;
    f23 *= delta;
    const double g1 = f13 * f22;
    const double g2 = f13 * f25 - f15 * f23;
    const double g3 = f11 * f23 - f13 * f21;
    const double g4 = -f13 * f24;
    const double g5 = f11 * f22;
    const double g6 = f11 * f25 - f15 * f21;
    const double g7 = -f15 * f24;
    const std::array<double, 5> coeffs =
    {
        {
            g5 * g5 + g1 * g1 + g3 * g3,
            2 * (g5 * g6 + g1 * g2 + g3 * g4),
            g6 * g6 + 2 * g5 * g7 + g2 * g2 + g4 * g4 - g1 * g1 - g3 * g3,
            2 * (g6 * g7 - g1 * g2 - g3 * g4),
            g7 * g7 - g2 * g2 - g4 * g4
        }
    };

    std::array<double, 4> s;
    solveQuarticPolynomial(coeffs, s);
    polishQuarticPolynomialRoots(coeffs, s);

    const Vec3 temp = k1.cross(nl);

    Mat3 Ck1nl;
    Ck1nl << k1, nl, temp;

    Mat3 Cb1k3tzT;
    Cb1k3tzT << b1.transpose(), k3.transpose(), tz.transpose();

    const Vec3 b3p = b3 * (delta / k3b3);

    for (const auto ctheta1p : s)
    {
        if (std::abs(ctheta1p) > 1)
            continue;
        const double stheta1p = ((k3b3 > 0) ? 1 : -1) * sqrt(1 - ctheta1p * ctheta1p);
        const double ntheta3 = stheta1p / ((g5 * ctheta1p + g6) * ctheta1p + g7);
        const double ctheta3 = (g1 * ctheta1p + g2) * ntheta3;
        const double stheta3 = (g3 * ctheta1p + g4) * ntheta3;

        Mat3 C13;
        C13 <<
            ctheta3,            0,         -stheta3,
                                stheta1p * stheta3, ctheta1p,  stheta1p * ctheta3,
                                ctheta1p * stheta3, -stheta1p, ctheta1p * ctheta3;

        const Mat3 R = (Ck1nl * C13) * Cb1k3tzT;
        const Vec3 rp3 = R.transpose() * w3; // R' * p3
        rotation_translation_solutions.emplace_back(R.transpose(), (b3p * stheta1p) - rp3);
    }
    return !rotation_translation_solutions.empty();
}

#define SUM_OR_DYNAMIC(x,y) (x==Eigen::Dynamic||y==Eigen::Dynamic)?Eigen::Dynamic:(x+y)

template<typename Derived1, typename Derived2>
struct hstack_return
{
    using Scalar =typename Derived1::Scalar;
    enum
    {
        RowsAtCompileTime = Derived1::RowsAtCompileTime,
        ColsAtCompileTime = SUM_OR_DYNAMIC( Derived1::ColsAtCompileTime, Derived2::ColsAtCompileTime ),
        Options = Derived1::Flags & Eigen::RowMajorBit ? Eigen::RowMajor : 0,
        MaxRowsAtCompileTime = Derived1::MaxRowsAtCompileTime,
        MaxColsAtCompileTime = SUM_OR_DYNAMIC( Derived1::MaxColsAtCompileTime, Derived2::MaxColsAtCompileTime )
    };
    using type =
        Eigen::Matrix<
        Scalar,
        RowsAtCompileTime,
        ColsAtCompileTime,
        Options,
        MaxRowsAtCompileTime,
        MaxColsAtCompileTime>;
};

template<typename Derived1, typename Derived2>
typename hstack_return<Derived1, Derived2>::type
HStack ( const Eigen::MatrixBase<Derived1>& lhs, const Eigen::MatrixBase<Derived2>& rhs )
{
    typename hstack_return<Derived1, Derived2>::type res;
    res.resize( lhs.rows(), lhs.cols() + rhs.cols() );
    res << lhs, rhs;
    return res;
}

std::vector<P3PEstimator_KE::M_t> P3PEstimator_KE::Estimate(const std::vector<P3PEstimator_KE::X_t>& points2D, const std::vector<P3PEstimator_KE::Y_t>& points3D)
{
    std::vector<std::tuple<Mat3, Vec3>> rotation_translation_solutions;
    std::vector<P3PEstimator_KE::M_t> models;
    if (computePoses(points2D, points3D, rotation_translation_solutions))
    {
        for (const auto & rot_trans_it : rotation_translation_solutions)
        {
            P3PEstimator_KE::M_t P;
            P = HStack(std::get<0>(rot_trans_it),std::get<1>(rot_trans_it));
            models.push_back(P);
        }
    }
    return models;
}

void P3PEstimator_KE::Residuals(const std::vector<X_t>& points2D,
                                const std::vector<Y_t>& points3D,
                                const M_t& proj_matrix,
                                std::vector<double>* residuals)
{
    ComputeSquaredReprojectionError(points2D, points3D, proj_matrix, residuals);
}

}  // namespace colmap
